// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/array/builder_nested.h>

#include <cstddef>
#include <cstdint>
#include <utility>
#include <vector>

#include <nebula/core/buffer.h>
#include <turbo/utility/status.h>
#include <nebula/types/type.h>
#include <turbo/base/checked_cast.h>
#include <turbo/log/logging.h>

namespace nebula {

    // ----------------------------------------------------------------------
    // VarLengthListLikeBuilder / BaseListBuilder / BaseListViewBuilder

    template
    class VarLengthListLikeBuilder<ListType>;

    template
    class VarLengthListLikeBuilder<LargeListType>;

    template
    class VarLengthListLikeBuilder<ListViewType>;

    template
    class VarLengthListLikeBuilder<LargeListViewType>;

    template
    class BaseListBuilder<ListType>;

    template
    class BaseListBuilder<LargeListType>;

    template
    class BaseListViewBuilder<ListViewType>;

    template
    class BaseListViewBuilder<LargeListViewType>;

    // ----------------------------------------------------------------------
    // MapBuilder

    MapBuilder::MapBuilder(MemoryPool *pool, const std::shared_ptr<ArrayBuilder> &key_builder,
                           std::shared_ptr<ArrayBuilder> const &item_builder,
                           const std::shared_ptr<DataType> &type)
            : ArrayBuilder(pool), key_builder_(key_builder), item_builder_(item_builder) {
        auto map_type = turbo::checked_cast<const MapType *>(type.get());
        entries_name_ = map_type->field(0)->name();
        key_name_ = map_type->key_field()->name();
        item_name_ = map_type->item_field()->name();
        item_nullable_ = map_type->item_field()->nullable();
        keys_sorted_ = map_type->keys_sorted();

        std::vector<std::shared_ptr<ArrayBuilder>> child_builders{key_builder, item_builder};
        auto struct_builder =
                std::make_shared<StructBuilder>(map_type->get_value_type(), pool, child_builders);

        list_builder_ =
                std::make_shared<ListBuilder>(pool, struct_builder, struct_builder->type());
    }

    MapBuilder::MapBuilder(MemoryPool *pool, const std::shared_ptr<ArrayBuilder> &key_builder,
                           const std::shared_ptr<ArrayBuilder> &item_builder,
                           bool keys_sorted)
            : MapBuilder(pool, key_builder, item_builder,
                         map(key_builder->type(), item_builder->type(), keys_sorted)) {}

    MapBuilder::MapBuilder(MemoryPool *pool,
                           const std::shared_ptr<ArrayBuilder> &struct_builder,
                           const std::shared_ptr<DataType> &type)
            : ArrayBuilder(pool) {
        auto map_type = turbo::checked_cast<const MapType *>(type.get());
        entries_name_ = map_type->field(0)->name();
        key_name_ = map_type->key_field()->name();
        item_name_ = map_type->item_field()->name();
        item_nullable_ = map_type->item_field()->nullable();
        keys_sorted_ = map_type->keys_sorted();
        key_builder_ = struct_builder->child_builder(0);
        item_builder_ = struct_builder->child_builder(1);
        list_builder_ =
                std::make_shared<ListBuilder>(pool, struct_builder, struct_builder->type());
    }

    turbo::Status MapBuilder::resize(int64_t capacity) {
        TURBO_RETURN_NOT_OK(list_builder_->resize(capacity));
        capacity_ = list_builder_->capacity();
        return turbo::OkStatus();
    }

    void MapBuilder::reset() {
        list_builder_->reset();
        ArrayBuilder::reset();
    }

    turbo::Status MapBuilder::finish_internal(std::shared_ptr<ArrayData> *out) {
        KCHECK_EQ(item_builder_->length(), key_builder_->length())
        << "keys and items builders don't have the same size in MapBuilder";
        TURBO_RETURN_NOT_OK(AdjustStructBuilderLength());
        TURBO_RETURN_NOT_OK(list_builder_->finish_internal(out));
        (*out)->type = type();
        ArrayBuilder::reset();
        return turbo::OkStatus();
    }

    turbo::Status MapBuilder::append_values(const int32_t *offsets, int64_t length,
                                           const uint8_t *valid_bytes) {
                DKCHECK_EQ(item_builder_->length(), key_builder_->length());
        TURBO_RETURN_NOT_OK(AdjustStructBuilderLength());
        TURBO_RETURN_NOT_OK(list_builder_->append_values(offsets, length, valid_bytes));
        length_ = list_builder_->length();
        null_count_ = list_builder_->null_count();
        return turbo::OkStatus();
    }

    turbo::Status MapBuilder::append() {
                DKCHECK_EQ(item_builder_->length(), key_builder_->length());
        TURBO_RETURN_NOT_OK(AdjustStructBuilderLength());
        TURBO_RETURN_NOT_OK(list_builder_->append());
        length_ = list_builder_->length();
        return turbo::OkStatus();
    }

    turbo::Status MapBuilder::append_null() {
                DKCHECK_EQ(item_builder_->length(), key_builder_->length());
        TURBO_RETURN_NOT_OK(AdjustStructBuilderLength());
        TURBO_RETURN_NOT_OK(list_builder_->append_null());
        length_ = list_builder_->length();
        null_count_ = list_builder_->null_count();
        return turbo::OkStatus();
    }

    turbo::Status MapBuilder::append_nulls(int64_t length) {
                DKCHECK_EQ(item_builder_->length(), key_builder_->length());
        TURBO_RETURN_NOT_OK(AdjustStructBuilderLength());
        TURBO_RETURN_NOT_OK(list_builder_->append_nulls(length));
        length_ = list_builder_->length();
        null_count_ = list_builder_->null_count();
        return turbo::OkStatus();
    }

    turbo::Status MapBuilder::append_empty_value() {
                DKCHECK_EQ(item_builder_->length(), key_builder_->length());
        TURBO_RETURN_NOT_OK(AdjustStructBuilderLength());
        TURBO_RETURN_NOT_OK(list_builder_->append_empty_value());
        length_ = list_builder_->length();
        null_count_ = list_builder_->null_count();
        return turbo::OkStatus();
    }

    turbo::Status MapBuilder::append_empty_values(int64_t length) {
                DKCHECK_EQ(item_builder_->length(), key_builder_->length());
        TURBO_RETURN_NOT_OK(AdjustStructBuilderLength());
        TURBO_RETURN_NOT_OK(list_builder_->append_empty_values(length));
        length_ = list_builder_->length();
        null_count_ = list_builder_->null_count();
        return turbo::OkStatus();
    }

    turbo::Status MapBuilder::AdjustStructBuilderLength() {
        // If key/item builders have been appended, adjust struct builder length
        // to match. Struct and key are non-nullable, append all valid values.
        auto struct_builder =
                turbo::checked_cast<StructBuilder *>(list_builder_->value_builder());
        if (struct_builder->length() < key_builder_->length()) {
            int64_t length_diff = key_builder_->length() - struct_builder->length();
            TURBO_RETURN_NOT_OK(struct_builder->append_values(length_diff, nullptr));
        }
        return turbo::OkStatus();
    }

    // ----------------------------------------------------------------------
    // FixedSizeListBuilder

    FixedSizeListBuilder::FixedSizeListBuilder(
            MemoryPool *pool, const std::shared_ptr<ArrayBuilder> &value_builder,
            const std::shared_ptr<DataType> &type)
            : ArrayBuilder(pool),
              value_field_(type->field(0)),
              list_size_(
                      turbo::checked_cast<const FixedSizeListType *>(type.get())->list_size()),
              value_builder_(value_builder) {}

    FixedSizeListBuilder::FixedSizeListBuilder(
            MemoryPool *pool, const std::shared_ptr<ArrayBuilder> &value_builder,
            int32_t list_size)
            : FixedSizeListBuilder(pool, value_builder,
                                   fixed_size_list(value_builder->type(), list_size)) {}

    void FixedSizeListBuilder::reset() {
        ArrayBuilder::reset();
        value_builder_->reset();
    }

    turbo::Status FixedSizeListBuilder::append() {
        TURBO_RETURN_NOT_OK(Reserve(1));
        unsafe_append_to_bitmap(true);
        return turbo::OkStatus();
    }

    turbo::Status FixedSizeListBuilder::append_values(int64_t length, const uint8_t *valid_bytes) {
        TURBO_RETURN_NOT_OK(Reserve(length));
        unsafe_append_to_bitmap(valid_bytes, length);
        return turbo::OkStatus();
    }

    turbo::Status FixedSizeListBuilder::append_null() {
        TURBO_RETURN_NOT_OK(Reserve(1));
        unsafe_append_to_bitmap(false);
        return value_builder_->append_nulls(list_size_);
    }

    turbo::Status FixedSizeListBuilder::append_nulls(int64_t length) {
        TURBO_RETURN_NOT_OK(Reserve(length));
        unsafe_append_to_bitmap(length, false);
        return value_builder_->append_nulls(list_size_ * length);
    }

    turbo::Status FixedSizeListBuilder::validate_overflow(int64_t new_elements) {
        auto new_length = value_builder_->length() + new_elements;
        if (new_elements != list_size_) {
            return turbo::invalid_argument_error("Length of item not correct: expected ", list_size_,
                                                 " but got array of size ", new_elements);
        }
        if (new_length > maximum_elements()) {
            return turbo::resource_exhausted_error("array cannot contain more than ", maximum_elements(),
                                                   " elements, have ", new_elements);
        }
        return turbo::OkStatus();
    }

    turbo::Status FixedSizeListBuilder::append_empty_value() {
        TURBO_RETURN_NOT_OK(Reserve(1));
        unsafe_append_to_bitmap(true);
        return value_builder_->append_empty_values(list_size_);
    }

    turbo::Status FixedSizeListBuilder::append_empty_values(int64_t length) {
        TURBO_RETURN_NOT_OK(Reserve(length));
        unsafe_append_to_bitmap(length, true);
        return value_builder_->append_empty_values(list_size_ * length);
    }

    turbo::Status FixedSizeListBuilder::resize(int64_t capacity) {
        TURBO_RETURN_NOT_OK(check_capacity(capacity));
        return ArrayBuilder::resize(capacity);
    }

    turbo::Status FixedSizeListBuilder::finish_internal(std::shared_ptr<ArrayData> *out) {
        std::shared_ptr<ArrayData> items;

        if (value_builder_->length() == 0) {
            // Try to make sure we get a non-null values buffer (ARROW-2744)
            TURBO_RETURN_NOT_OK(value_builder_->resize(0));
        }
        TURBO_RETURN_NOT_OK(value_builder_->finish_internal(&items));

        std::shared_ptr<Buffer> null_bitmap;
        TURBO_RETURN_NOT_OK(null_bitmap_builder_.finish(&null_bitmap));
        *out = ArrayData::create(type(), length_, {null_bitmap}, {std::move(items)}, null_count_);
        reset();
        return turbo::OkStatus();
    }

    // ----------------------------------------------------------------------
    // Struct

    StructBuilder::StructBuilder(const std::shared_ptr<DataType> &type, MemoryPool *pool,
                                 std::vector<std::shared_ptr<ArrayBuilder>> field_builders)
            : ArrayBuilder(pool), type_(type) {
        children_ = std::move(field_builders);
    }

    void StructBuilder::reset() {
        ArrayBuilder::reset();
        for (const auto &field_builder: children_) {
            field_builder->reset();
        }
    }

    turbo::Status StructBuilder::finish_internal(std::shared_ptr<ArrayData> *out) {
        std::shared_ptr<Buffer> null_bitmap;
        TURBO_RETURN_NOT_OK(null_bitmap_builder_.finish(&null_bitmap));

        std::vector<std::shared_ptr<ArrayData>> child_data(children_.size());
        for (size_t i = 0; i < children_.size(); ++i) {
            if (length_ == 0) {
                // Try to make sure the child buffers are initialized
                TURBO_RETURN_NOT_OK(children_[i]->resize(0));
            }
            TURBO_RETURN_NOT_OK(children_[i]->finish_internal(&child_data[i]));
        }

        *out = ArrayData::create(type(), length_, {null_bitmap}, null_count_);
        (*out)->child_data = std::move(child_data);

        capacity_ = length_ = null_count_ = 0;
        return turbo::OkStatus();
    }

    std::shared_ptr<DataType> StructBuilder::type() const {
                DKCHECK_EQ(type_->fields().size(), children_.size());
        std::vector<std::shared_ptr<Field>> fields(children_.size());
        for (int i = 0; i < static_cast<int>(fields.size()); ++i) {
            fields[i] = type_->field(i)->with_type(children_[i]->type());
        }
        return STRUCT(std::move(fields));
    }

}  // namespace nebula
